#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/platform_device.h>
#include <linux/of.h>
#include <linux/of_device.h>          // 新增：用于 bus_find_device_by_of_node
#include <linux/dmaengine.h>
#include <linux/dma-mapping.h>
#include <linux/slab.h>
#include <linux/completion.h>
#include <linux/delay.h>
#include <linux/device.h>

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Your Name");
MODULE_DESCRIPTION("PL AXI DMA Memory Copy Module for Zynq (Linux 5.4 Compatible)");
MODULE_VERSION("1.0");

// DMA 通道
static struct dma_chan *dma_chan_tx = NULL;  // MM2S: 内存 -> PL
static struct dma_chan *dma_chan_rx = NULL;  // S2MM: PL -> 内存
static struct platform_device *dma_pdev = NULL;

// 缓冲区
static void *src_buffer = NULL;
static void *dst_buffer = NULL;
static dma_addr_t src_phys = 0;
static dma_addr_t dst_phys = 0;

#define COPY_SIZE 4096

// 同步机制
static struct completion tx_done;
static struct completion rx_done;
static int tx_error = 0;
static int rx_error = 0;

// ------------------------
// 正确的错误回调函数（Linux 5.4+）
// ------------------------

// 注意：新 API 使用 callback_result 接收状态
static void dma_tx_callback(void *param)
{
    pr_info("DMA TX (MM2S) completed\n");
    complete(&tx_done);
}

static void dma_rx_callback(void *param)
{
    pr_info("DMA RX (S2MM) completed\n");
    complete(&rx_done);
}

// 错误回调现在通过 callback_result 传递
static void dma_tx_callback_result(void *param, const struct dmaengine_result *result)
{
    if (result->result != DMA_TRANS_NOERROR) {
        pr_err("DMA TX error: %d\n", result->result);
        tx_error = -EIO;
    }
    complete(&tx_done);
}

static void dma_rx_callback_result(void *param, const struct dmaengine_result *result)
{
    if (result->result != DMA_TRANS_NOERROR) {
        pr_err("DMA RX error: %d\n", result->result);
        rx_error = -EIO;
    }
    complete(&rx_done);
}

// ------------------------
// 初始化函数（关键修改）
// ------------------------

static int __init pl_axi_dma_init(void)
{
    struct device_node *np;
    int ret = 0;

    pr_info("PL AXI DMA module initializing\n");

    // 1. 查找设备树节点
    np = of_find_node_by_path("/amba_pl/dma@40400000");
    if (!np) {
        pr_err("Failed to find device node '/amba_pl/dma@40400000'\n");
        return -ENODEV;
    }

    // 2. 替代 of_find_device_by_node() → 使用 bus_find_device_by_of_node()
    struct device *dev = bus_find_device_by_of_node(&platform_bus_type, np);
    of_node_put(np);

    if (!dev) {
        pr_err("Failed to find device from device tree node\n");
        return -ENODEV;
    }

    dma_pdev = to_platform_device(dev);
    pr_info("Found platform device: %s\n", dev_name(&dma_pdev->dev));

    // 3. 请求 TX 通道 (MM2S)
    dma_chan_tx = dma_request_chan(&dma_pdev->dev, "mm2s");
    if (IS_ERR(dma_chan_tx)) {
        ret = PTR_ERR(dma_chan_tx);
        pr_err("Failed to request TX (mm2s) channel: %d\n", ret);
        dma_chan_tx = NULL;
        goto err_put_device;
    }
    pr_info("Successfully requested TX channel: %s\n", dma_chan_name(dma_chan_tx));

    // 4. 请求 RX 通道 (S2MM)
    dma_chan_rx = dma_request_chan(&dma_pdev->dev, "s2mm");
    if (IS_ERR(dma_chan_rx)) {
        ret = PTR_ERR(dma_chan_rx);
        pr_err("Failed to request RX (s2mm) channel: %d\n", ret);
        dma_chan_rx = NULL;
        goto err_release_tx;
    }
    pr_info("Successfully requested RX channel: %s\n", dma_chan_name(dma_chan_rx));

    pr_info("PL AXI DMA channels initialized successfully!\n");

    // ========================
    // 分配内存（CMA）
    // ========================

    src_buffer = dma_alloc_coherent(&dma_pdev->dev, COPY_SIZE, &src_phys, GFP_KERNEL);
    if (!src_buffer) {
        pr_err("Failed to allocate source buffer\n");
        ret = -ENOMEM;
        goto err_release_rx;
    }

    dst_buffer = dma_alloc_coherent(&dma_pdev->dev, COPY_SIZE, &dst_phys, GFP_KERNEL);
    if (!dst_buffer) {
        pr_err("Failed to allocate destination buffer\n");
        ret = -ENOMEM;
        goto err_free_src;
    }

    pr_info("Allocated buffers: src=0x%p (%pad), dst=0x%p (%pad)\n",
            src_buffer, &src_phys, dst_buffer, &dst_phys);

    memset(src_buffer, 0, COPY_SIZE);
    strcpy(src_buffer, "Hello from AXI DMA! This is a test message for PL-to-PL memory copy via MM2S+S2MM.");
    memset(dst_buffer, 0, COPY_SIZE);

    init_completion(&tx_done);
    init_completion(&rx_done);
    tx_error = 0;
    rx_error = 0;

    // ========================
    // 提交 MM2S 传输（内存 → PL）
    // ========================
    struct dma_async_tx_descriptor *tx_desc;
    tx_desc = dmaengine_prep_slave_single(dma_chan_tx, src_phys,
                                          COPY_SIZE, DMA_MEM_TO_DEV,
                                          DMA_PREP_INTERRUPT | DMA_CTRL_ACK);
    if (!tx_desc) {
        pr_err("Failed to prepare TX descriptor\n");
        ret = -EIO;
        goto err_free_dst;
    }

    // 设置完成回调
    tx_desc->callback = dma_tx_callback;
    tx_desc->callback_param = NULL;

    // ✅ 设置错误/结果回调（新 API）
    tx_desc->callback_result = dma_tx_callback_result;
    tx_desc->callback_param = NULL;

    dmaengine_submit(tx_desc);
    dma_async_issue_pending(dma_chan_tx);

    pr_info("Submitted MM2S transfer: %zu bytes from %pad\n", COPY_SIZE, &src_phys);

    // ========================
    // 提交 S2MM 传输（PL → 内存）
    // ========================
    struct dma_async_tx_descriptor *rx_desc;
    rx_desc = dmaengine_prep_slave_single(dma_chan_rx, dst_phys,
                                          COPY_SIZE, DMA_DEV_TO_MEM,
                                          DMA_PREP_INTERRUPT | DMA_CTRL_ACK);
    if (!rx_desc) {
        pr_err("Failed to prepare RX descriptor\n");
        ret = -EIO;
        goto err_terminate_tx;
    }

    rx_desc->callback = dma_rx_callback;
    rx_desc->callback_param = NULL;

    rx_desc->callback_result = dma_rx_callback_result;
    rx_desc->callback_param = NULL;

    dmaengine_submit(rx_desc);
    dma_async_issue_pending(dma_chan_rx);

    pr_info("Submitted S2MM transfer: %zu bytes to %pad\n", COPY_SIZE, &dst_phys);

    // ========================
    // 等待完成
    // ========================
    if (wait_for_completion_timeout(&tx_done, msecs_to_jiffies(5000)) == 0) {
        pr_err("TX timeout!\n");
        ret = -ETIMEDOUT;
        goto err_terminate_all;
    }

    if (wait_for_completion_timeout(&rx_done, msecs_to_jiffies(5000)) == 0) {
        pr_err("RX timeout!\n");
        ret = -ETIMEDOUT;
        goto err_terminate_all;
    }

    if (tx_error || rx_error) {
        pr_err("DMA transfer failed with error\n");
        ret = -EIO;
        goto err_terminate_all;
    }

    pr_info("DMA transfer completed successfully!\n");

    // ========================
    // 验证结果
    // ========================
    size_t len = strlen((char *)src_buffer);
    if (memcmp(src_buffer, dst_buffer, len) == 0) {
        pr_info("✅ Data verification PASSED!\n");
        pr_info("Copied: %s\n", (char *)dst_buffer);
    } else {
        pr_err("❌ Data verification FAILED!\n");
        ret = -EIO;
    }

    return 0;

err_terminate_all:
    dmaengine_terminate_sync(dma_chan_tx);
    dmaengine_terminate_sync(dma_chan_rx);
err_terminate_tx:
    dmaengine_terminate_sync(dma_chan_tx);
err_free_dst:
    dma_free_coherent(&dma_pdev->dev, COPY_SIZE, dst_buffer, dst_phys);
    dst_buffer = NULL;
err_free_src:
    dma_free_coherent(&dma_pdev->dev, COPY_SIZE, src_buffer, src_phys);
    src_buffer = NULL;
err_release_rx:
    if (dma_chan_rx) {
        dma_release_channel(dma_chan_rx);
        dma_chan_rx = NULL;
    }
err_release_tx:
    if (dma_chan_tx) {
        dma_release_channel(dma_chan_tx);
        dma_chan_tx = NULL;
    }
err_put_device:
    if (dma_pdev) {
        put_device(&dma_pdev->dev);  // 对应 bus_find_device_by_of_node 的引用
    }
    return ret;
}

// ------------------------
// 退出函数
// ------------------------

static void __exit pl_axi_dma_exit(void)
{
    pr_info("Releasing PL AXI DMA channels and buffers\n");

    if (dma_chan_tx) {
        dmaengine_terminate_sync(dma_chan_tx);
        dma_release_channel(dma_chan_tx);
        dma_chan_tx = NULL;
    }

    if (dma_chan_rx) {
        dmaengine_terminate_sync(dma_chan_rx);
        dma_release_channel(dma_chan_rx);
        dma_chan_rx = NULL;
    }

    if (src_buffer) {
        dma_free_coherent(&dma_pdev->dev, COPY_SIZE, src_buffer, src_phys);
        src_buffer = NULL;
    }

    if (dst_buffer) {
        dma_free_coherent(&dma_pdev->dev, COPY_SIZE, dst_buffer, dst_phys);
        dst_buffer = NULL;
    }

    if (dma_pdev) {
        put_device(&dma_pdev->dev);  // 释放 device 引用
        dma_pdev = NULL;
    }

    pr_info("PL AXI DMA module exited\n");
}

module_init(pl_axi_dma_init);
module_exit(pl_axi_dma_exit);
